"""
RagFlow 命令行工具
提供命令行界面操作 RagFlow 客户端
"""

import argparse
import json
import sys
from typing import Dict, List, Any, Optional

from ragflow_client.config import RagFlowConfig
from ragflow_client.service import RagFlowService
from ragflow_client.utils.logger import get_logger

logger = get_logger(__name__)


def print_json(data: Any) -> None:
    """
    格式化输出 JSON 数据
    
    Args:
        data: 要输出的数据
    """
    print(json.dumps(data, ensure_ascii=False, indent=2))


def handle_datasets(args: argparse.Namespace, service: RagFlowService) -> int:
    """
    处理知识库相关命令
    
    Args:
        args: 命令行参数
        service: RagFlow 服务对象
        
    Returns:
        退出码，0 表示成功，非 0 表示失败
    """
    if args.list:
        # 列出知识库
        status, result = service.list_datasets(args.page, args.page_size)
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    elif args.get:
        # 获取知识库详情
        status, result = service.get_dataset_info(args.get)
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    elif args.create:
        # 创建知识库
        status, result = service.create_dataset(args.create, args.description, args.embedding_model)
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    elif args.update:
        # 更新知识库
        status, result = service.update_dataset(args.update, args.name, args.description)
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    elif args.delete:
        # 删除知识库
        status, result = service.delete_dataset(args.delete)
        if status:
            print(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    else:
        print("错误: 请指定知识库操作")
        return 1


def handle_chunks(args: argparse.Namespace, service: RagFlowService) -> int:
    """
    处理知识块相关命令
    
    Args:
        args: 命令行参数
        service: RagFlow 服务对象
        
    Returns:
        退出码，0 表示成功，非 0 表示失败
    """
    if args.list:
        # 列出知识块
        status, result = service.list_chunks(args.dataset_id, args.document_id, args.page, args.page_size)
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    elif args.get:
        # 获取知识块详情
        status, result = service.get_chunk_info(args.get, args.dataset_id)
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    elif args.create:
        # 创建知识块
        metadata = None
        if args.metadata:
            try:
                metadata = json.loads(args.metadata)
            except json.JSONDecodeError:
                print("错误: 元数据必须是有效的 JSON 格式")
                return 1
        
        status, result = service.create_chunk(args.create, args.dataset_id, metadata)
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    elif args.update:
        # 更新知识块
        metadata = None
        if args.metadata:
            try:
                metadata = json.loads(args.metadata)
            except json.JSONDecodeError:
                print("错误: 元数据必须是有效的 JSON 格式")
                return 1
        
        status, result = service.update_chunk(args.update, args.content, args.dataset_id, metadata)
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    elif args.delete:
        # 删除知识块
        status, result = service.delete_chunk(args.delete, args.dataset_id)
        if status:
            print(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    else:
        print("错误: 请指定知识块操作")
        return 1


def handle_documents(args: argparse.Namespace, service: RagFlowService) -> int:
    """
    处理文档相关命令
    
    Args:
        args: 命令行参数
        service: RagFlow 服务对象
        
    Returns:
        退出码，0 表示成功，非 0 表示失败
    """
    if args.list:
        # 列出文档
        status, result = service.list_documents(args.dataset_id, args.page, args.page_size)
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    elif args.get:
        # 获取文档详情
        status, result = service.get_document_info(args.get, args.dataset_id)
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    elif args.upload:
        # 上传文档
        status, result = service.api.upload_files(args.upload, args.dataset_id)
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    elif args.parse:
        # 解析文档
        status, result = service.parse_document(args.parse, args.dataset_id, args.parser_id)
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    elif args.delete:
        # 删除文档
        status, result = service.delete_document(args.delete, args.dataset_id)
        if status:
            print(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    else:
        print("错误: 请指定文档操作")
        return 1


def handle_search(args: argparse.Namespace, service: RagFlowService) -> int:
    """
    处理搜索命令
    
    Args:
        args: 命令行参数
        service: RagFlow 服务对象
        
    Returns:
        退出码，0 表示成功，非 0 表示失败
    """
    status, result = service.search(args.query, args.dataset_id, args.top_k, args.filter_threshold)
    if status:
        print_json(result)
        return 0
    else:
        print(f"错误: {result}")
        return 1


def handle_chat(args: argparse.Namespace, service: RagFlowService) -> int:
    """
    处理对话命令
    
    Args:
        args: 命令行参数
        service: RagFlow 服务对象
        
    Returns:
        退出码，0 表示成功，非 0 表示失败
    """
    history = None
    if args.history:
        try:
            with open(args.history, 'r', encoding='utf-8') as f:
                history = json.load(f)
        except Exception as e:
            print(f"错误: 无法读取对话历史文件: {str(e)}")
            return 1
    
    status, result = service.chat(args.query, args.dataset_id, history, args.stream, args.top_k, args.filter_threshold)
    
    if status:
        if args.stream:
            # 处理流式响应
            try:
                for chunk in result.iter_lines():
                    if chunk:
                        chunk_data = json.loads(chunk.decode('utf-8'))
                        if 'data' in chunk_data:
                            content = chunk_data['data'].get('content', '')
                            print(content, end='', flush=True)
                print()  # 换行
                return 0
            except Exception as e:
                print(f"\n错误: 处理流式响应失败: {str(e)}")
                return 1
        else:
            # 处理普通响应
            print(result.get('answer', ''))
            
            # 如果需要输出完整响应
            if args.verbose:
                print("\n完整响应:")
                print_json(result)
            
            return 0
    else:
        print(f"错误: {result}")
        return 1


def handle_system(args: argparse.Namespace, service: RagFlowService) -> int:
    """
    处理系统信息命令
    
    Args:
        args: 命令行参数
        service: RagFlow 服务对象
        
    Returns:
        退出码，0 表示成功，非 0 表示失败
    """
    if args.parsers:
        # 获取解析器列表
        status, result = service.get_parsers()
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    elif args.embedding_models:
        # 获取嵌入模型列表
        status, result = service.get_embedding_models()
        if status:
            print_json(result)
            return 0
        else:
            print(f"错误: {result}")
            return 1
    
    else:
        print("错误: 请指定系统信息操作")
        return 1


def main(args: Optional[List[str]] = None) -> int:
    """
    命令行工具主函数
    
    Args:
        args: 命令行参数，如果为 None 则从 sys.argv 获取
        
    Returns:
        退出码，0 表示成功，非 0 表示失败
    """
    parser = argparse.ArgumentParser(description="RagFlow 命令行工具")
    subparsers = parser.add_subparsers(dest="command", help="子命令")
    
    # 知识库命令
    dataset_parser = subparsers.add_parser("dataset", help="知识库操作")
    dataset_group = dataset_parser.add_mutually_exclusive_group(required=True)
    dataset_group.add_argument("--list", action="store_true", help="列出知识库")
    dataset_group.add_argument("--get", metavar="ID", help="获取知识库详情")
    dataset_group.add_argument("--create", metavar="NAME", help="创建知识库")
    dataset_group.add_argument("--update", metavar="ID", help="更新知识库")
    dataset_group.add_argument("--delete", metavar="ID", help="删除知识库")
    dataset_parser.add_argument("--name", help="知识库名称，用于更新")
    dataset_parser.add_argument("--description", help="知识库描述")
    dataset_parser.add_argument("--embedding-model", default="text2vec-base", help="嵌入模型名称")
    dataset_parser.add_argument("--page", type=int, default=1, help="页码")
    dataset_parser.add_argument("--page-size", type=int, default=10, help="每页数量")
    
    # 文档命令
    document_parser = subparsers.add_parser("document", help="文档操作")
    document_group = document_parser.add_mutually_exclusive_group(required=True)
    document_group.add_argument("--list", action="store_true", help="列出文档")
    document_group.add_argument("--get", metavar="ID", help="获取文档详情")
    document_group.add_argument("--upload", metavar="FILE", nargs="+", help="上传文档")
    document_group.add_argument("--parse", metavar="ID", help="解析文档")
    document_group.add_argument("--delete", metavar="ID", help="删除文档")
    document_parser.add_argument("--dataset-id", help="知识库 ID")
    document_parser.add_argument("--parser-id", help="解析器 ID")
    document_parser.add_argument("--page", type=int, default=1, help="页码")
    document_parser.add_argument("--page-size", type=int, default=10, help="每页数量")
    
    # 知识块命令
    chunk_parser = subparsers.add_parser("chunk", help="知识块操作")
    chunk_group = chunk_parser.add_mutually_exclusive_group(required=True)
    chunk_group.add_argument("--list", action="store_true", help="列出知识块")
    chunk_group.add_argument("--get", metavar="ID", help="获取知识块详情")
    chunk_group.add_argument("--create", metavar="CONTENT", help="创建知识块")
    chunk_group.add_argument("--update", metavar="ID", help="更新知识块")
    chunk_group.add_argument("--delete", metavar="ID", help="删除知识块")
    chunk_parser.add_argument("--dataset-id", help="知识库 ID")
    chunk_parser.add_argument("--document-id", help="文档 ID，用于列出特定文档的知识块")
    chunk_parser.add_argument("--content", help="知识块内容，用于更新")
    chunk_parser.add_argument("--metadata", help="知识块元数据，JSON 格式")
    chunk_parser.add_argument("--page", type=int, default=1, help="页码")
    chunk_parser.add_argument("--page-size", type=int, default=10, help="每页数量")
    
    # 搜索命令
    search_parser = subparsers.add_parser("search", help="搜索知识库")
    search_parser.add_argument("query", help="查询文本")
    search_parser.add_argument("--dataset-id", help="知识库 ID")
    search_parser.add_argument("--top-k", type=int, default=5, help="返回结果数量")
    search_parser.add_argument("--filter-threshold", type=float, default=0.0, help="过滤阈值")
    
    # 对话命令
    chat_parser = subparsers.add_parser("chat", help="与知识库对话")
    chat_parser.add_argument("query", help="用户问题")
    chat_parser.add_argument("--dataset-id", help="知识库 ID")
    chat_parser.add_argument("--history", help="对话历史文件路径（JSON 格式）")
    chat_parser.add_argument("--stream", action="store_true", help="使用流式响应")
    chat_parser.add_argument("--top-k", type=int, default=5, help="检索结果数量")
    chat_parser.add_argument("--filter-threshold", type=float, default=0.0, help="过滤阈值")
    chat_parser.add_argument("--verbose", action="store_true", help="输出完整响应")
    
    # 系统信息命令
    system_parser = subparsers.add_parser("system", help="系统信息")
    system_group = system_parser.add_mutually_exclusive_group(required=True)
    system_group.add_argument("--parsers", action="store_true", help="获取解析器列表")
    system_group.add_argument("--embedding-models", action="store_true", help="获取嵌入模型列表")
    
    # 解析命令行参数
    parsed_args = parser.parse_args(args)
    
    if not parsed_args.command:
        parser.print_help()
        return 1
    
    try:
        # 加载配置
        config = RagFlowConfig.from_env()
        
        # 创建服务
        service = RagFlowService(config)
        
        # 处理命令
        if parsed_args.command == "dataset":
            return handle_datasets(parsed_args, service)
        elif parsed_args.command == "document":
            return handle_documents(parsed_args, service)
        elif parsed_args.command == "chunk":
            return handle_chunks(parsed_args, service)
        elif parsed_args.command == "search":
            return handle_search(parsed_args, service)
        elif parsed_args.command == "chat":
            return handle_chat(parsed_args, service)
        elif parsed_args.command == "system":
            return handle_system(parsed_args, service)
        else:
            print(f"错误: 未知命令 {parsed_args.command}")
            return 1
            
    except Exception as e:
        print(f"错误: {str(e)}")
        return 1


if __name__ == "__main__":
    sys.exit(main())